// Copyright 2014 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "media/filters/h264_bitstream_buffer.h"

#include "base/sys_byteorder.h"

namespace media {

H264BitstreamBuffer::H264BitstreamBuffer()
    : data_(NULL)
{
    Reset();
}

H264BitstreamBuffer::~H264BitstreamBuffer()
{
    free(data_);
    data_ = NULL;
}

void H264BitstreamBuffer::Reset()
{
    free(data_);
    data_ = NULL;

    capacity_ = 0;
    pos_ = 0;
    reg_ = 0;

    Grow();

    bits_left_in_reg_ = kRegBitSize;
}

void H264BitstreamBuffer::Grow()
{
    data_ = static_cast<uint8_t*>(realloc(data_, capacity_ + kGrowBytes));
    CHECK(data_) << "Failed growing the buffer";
    capacity_ += kGrowBytes;
}

void H264BitstreamBuffer::FlushReg()
{
    // Flush all bytes that have at least one bit cached, but not more
    // (on Flush(), reg_ may not be full).
    size_t bits_in_reg = kRegBitSize - bits_left_in_reg_;
    if (bits_in_reg == 0)
        return;

    size_t bytes_in_reg = (bits_in_reg + 7) / 8;
    reg_ <<= (kRegBitSize - bits_in_reg);

    // Convert to MSB and append as such to the stream.
    reg_ = base::HostToNet64(reg_);

    // Make sure we have enough space. Grow() will CHECK() on allocation failure.
    if (pos_ + bytes_in_reg < capacity_)
        Grow();

    memcpy(data_ + pos_, &reg_, bytes_in_reg);
    pos_ += bytes_in_reg;

    reg_ = 0;
    bits_left_in_reg_ = kRegBitSize;
}

void H264BitstreamBuffer::AppendU64(size_t num_bits, uint64_t val)
{
    CHECK_LE(num_bits, kRegBitSize);

    while (num_bits > 0) {
        if (bits_left_in_reg_ == 0)
            FlushReg();

        uint64_t bits_to_write = num_bits > bits_left_in_reg_ ? bits_left_in_reg_ : num_bits;
        uint64_t val_to_write = (val >> (num_bits - bits_to_write));
        if (bits_to_write < 64)
            val_to_write &= ((1ull << bits_to_write) - 1);
        reg_ <<= bits_to_write;
        reg_ |= val_to_write;
        num_bits -= bits_to_write;
        bits_left_in_reg_ -= bits_to_write;
    }
}

void H264BitstreamBuffer::AppendBool(bool val)
{
    if (bits_left_in_reg_ == 0)
        FlushReg();

    reg_ <<= 1;
    reg_ |= (static_cast<uint64_t>(val) & 1);
    --bits_left_in_reg_;
}

void H264BitstreamBuffer::AppendSE(int val)
{
    if (val > 0)
        AppendUE(val * 2 - 1);
    else
        AppendUE(-val * 2);
}

void H264BitstreamBuffer::AppendUE(unsigned int val)
{
    size_t num_zeros = 0;
    unsigned int v = val + 1;

    while (v > 1) {
        v >>= 1;
        ++num_zeros;
    }

    AppendBits(num_zeros, 0);
    AppendBits(num_zeros + 1, val + 1);
}

#define DCHECK_FINISHED()                                                        \
    DCHECK_EQ(bits_left_in_reg_, kRegBitSize) << "Pending bits not yet written " \
                                                 "to the buffer, call "          \
                                                 "FinishNALU() first."

void H264BitstreamBuffer::BeginNALU(H264NALU::Type nalu_type, int nal_ref_idc)
{
    DCHECK_FINISHED();

    DCHECK_LE(nalu_type, H264NALU::kEOStream);
    DCHECK_GE(nal_ref_idc, 0);
    DCHECK_LE(nal_ref_idc, 3);

    AppendBits(32, 0x00000001);
    AppendBits(1, 0); // forbidden_zero_bit
    AppendBits(2, nal_ref_idc);
    AppendBits(5, nalu_type);
}

void H264BitstreamBuffer::FinishNALU()
{
    // RBSP stop one bit.
    AppendBits(1, 1);

    // Byte-alignment zero bits.
    AppendBits(bits_left_in_reg_ % 8, 0);

    if (bits_left_in_reg_ != kRegBitSize)
        FlushReg();
}

size_t H264BitstreamBuffer::BytesInBuffer()
{
    DCHECK_FINISHED();
    return pos_;
}

uint8_t* H264BitstreamBuffer::data()
{
    DCHECK(data_);
    DCHECK_FINISHED();

    return data_;
}

} // namespace media
